
import triton
import triton.language as tl

@triton.jit
def softmax(Y, y_stride_row, y_stride_col, X, x_stride_row, x_stride_col, R, C, debug):
    row_id = tl.program_id(0)
    BLOCK_SIZE=1024
    n = tl.arange(0, BLOCK_SIZE)
    cur_debug = debug + row_id * y_stride_row
    tl.store(cur_debug, row_id, mask=(n < C))

import torch
# import torch_npu
X = torch.normal(0, 1, size=(583, 931), device='cpu:0')
Y = torch.empty_like(X)
debug = torch.empty_like(X)
grid = (X.shape[0], )
softmax[grid](
    Y, Y.stride(0), Y.stride(1), 
    X, X.stride(0), X.stride(1), 
    X.shape[0], X.shape[1],
    debug
)
print(debug)
